'''
定义了包基类
'''
from scapy.all import (
	send, sendp, ls, sr, srp
)
from error import *

class BasicPacket(object):
	'''
	包基类，是一个迭代器，并且重载了加法、乘法运算符
	'''
	#魔法函数
	def __init__(self, *args:tuple, **kwargs:dict):
		'''
		param args: 初始化的包
		param layernum: 包所在的网络层数
		'''
		#初始化包
		if args:
			self.__packets__:list = list(args)
		else:
			self.__packets__:list = []
		self.__layernum__:int = kwargs.get('layernum', 0)

		#Lambda函数定义方法
		self.pop = lambda :self.__packets__.pop()
		self.length = lambda :len(self.__packets__)

	def __iter__(self):
		'''
		迭代初始化
		'''
		self.__index__:int = 0
		return self

	def __next__(self):
		'''
		每次迭代时调用
		'''
		if self.__packets__:
			if 0<=self.__index__<self.len():
				output = self.__packets__[self.__index__]
				self.__index__ += 1
				return output
		raise StopIteration()

	def __getitem__(self, index:int):
		'''
		下标查询
		'''
		return self.__packets__[index]

	def __str__(self) -> str:
		'''
		str类转换调用的函数
		'''
		return '\n'.join(self.summary())

	def __add__(self, packet):
		'''
		重载加法运算符
		'''
		newpacket = Packet(self.__packets__)
		newpacket.put(packet.__packet__)
		return newpacket

	def __mul__(self, num:int):
		'''
		重载乘法运算符
		'''
		self.__packets__ * num

	def __sub__(self, packet):
		'''
		重载减法运算符
		'''
		pass

	def packet(self, data:bytes):
		'''
		返回组装好的包
		'''
		if self.length==1:
			return self[0] if not data else self[0] / data
		else:
			output = self[0]
			for packet in self[1:]:
				output = output / packet
			return output if not data else output / data

	def put(self, *args) -> None:
		'''
		将包放入包列表中
		'''
		for inter_arg in args:
			try:
				for packet in args:
					break
			except:
				self.__packets__.append(inter_arg)
			else:
				self.__packets__ += inter_arg
	
	def clear(self):
		'''
		清除包列表中的包
		'''
		self.__packets__.clear()

	def ls(self, *args) -> None:
		'''
		使用ls函数分析包列表中所有包
		'''
		packets = args if args else self.__packets__
		for packet in packets:
			ls(packet)

	def summary(self) -> list[str]:
		'''
		返回所有包的summary结果
		'''
		return (packet.summary() for packet in self.__packets__)

	def prepare(self, **kwargs):
		return self.packet(
			kwargs.get('data', False)
		)

	def send(self, amount:int=1, **kwargs):
		'''
		发送数据包
		'''
		if self.__layernum__==2:
			def _send(packet, kwargs):
				return sendp(
						packet,
						verbose=kwargs.get('verbose', False),
						iface_hint=kwargs.get('iface_hint', None),
						return_packets=kwargs.get('return_packets', False),
						inter=kwargs.get('inter', 0),
						iface=kwargs.get('iface', None),
						socket=kwargs.get('socket', None),
						loop=kwargs.get('loop', 0),
						count=kwargs.get('count', None),
						realtime=kwargs.get('realtime', False)
				)
		elif self.__layernum__==3:
			def _send(packet, kwargs):
				return send(
						packet,
						verbose=kwargs.get('verbose', False),
						return_packets=kwargs.get('return_packets', False),
						inter=kwargs.get('inter', 0),
						iface=kwargs.get('iface', None),
						socket=kwargs.get('socket', None),
						loop=kwargs.get('loop', 0),
						count=kwargs.get('count', None),
						realtime=kwargs.get('realtime', False)
					)
			amount += 1
		else:
			raise SendTypeError('未知发送类型')
		return _send(
			(self.prepare(kwargs) for i in range(amount)), kwargs
		)

	def sendrecv(self, amount:int=1, **kwargs):
		'''
		发送和接收数据包
		'''
		if self.__layernum__==2:
			def _send(packet, kwargs):
				return sr(
					packet,
					promisc=kwargs.get('promisc', None),
					filter=kwargs.get('filter', None),
					iface=kwargs.get('iface', None),
					nofilter=kwargs.get('nofilter', 0),
					retry=kwargs.get('retry', 0),
					multi=kwargs.get('multi', False),
					rcv_pks=kwargs.get('rcv_pks', None),
					prebuild=kwargs.get('prebuild', False),
					threaded=kwargs.get('threades', False),
					session=kwargs.get('session', None)
				)
		elif self.__layernum__==3:
			def _send(packet, kwargs):
				return srp(
					packet,
					promisc=kwargs.get('promisc', None),
					filter=kwargs.get('filter', None),
					iface=kwargs.get('iface', None),
					nofilter=kwargs.get('nofilter', 0),
					retry=kwargs.get('retry', 0),
					multi=kwargs.get('multi', False),
					rcv_pks=kwargs.get('rcv_pks', None),
					prebuild=kwargs.get('prebuild', False),
					threaded=kwargs.get('threades', False),
					session=kwargs.get('session', None)
				)

	def search(self, protocol:str):
		'''
		从保存的包中搜索含有对应
		'''
		output:list = []
		protocol:str = protocol.upper()
		if protocol=='ETHER':
			protocol = 'Ether'
		for packet in self:
			try:
				packet[protocol]
			except:
				pass
			else:
				output.append(packet)
		return output
	
	def pick(self, protocol:str):
		'''
		从保存的包中分离出指定协议的包
		'''
		output:list = []
		protocol = protocol.upper()
		if protocol=='ETHER':
			protocol = 'Ether'
		_pick = lambda packet, protocol:packet[protocol]
		for packet in self:
			output.append(_pick(packet, protocol))
		return output